package cn.doitedu.ml.demo

import org.apache.spark.ml.classification.NaiveBayesModel
import org.apache.spark.ml.linalg.Vectors
import org.apache.spark.ml.linalg.Vector
import org.apache.spark.sql.SparkSession

import scala.collection.mutable

object NaiveBayesDemoTester {

  def main(args: Array[String]): Unit = {

    val spark = SparkSession.builder().appName("朴素贝叶斯分类算法示例：出轨预测").master("local").getOrCreate()
    import spark.implicits._
    import org.apache.spark.sql.functions._


    // 加载样本数据集
    val test = spark.read.option("header","true").csv("user_portrait/data/chugui/test")

    /**
     * 特征工程
     */
    // 特征值数字化
    val testDatumn = test.selectExpr(
      "name",
      "case when job='老师' then 1.0 when job='公务员' then 2.0   else 3.0 end as job",
      "case when income='低' then 1.0 when income='中' then 2.0   else 3.0 end as income",
      "case when age='中年' then 1.0 when age='青年' then 2.0   else 3.0 end as age",
      "case when sex='男' then 1.0  else 2.0 end as sex"
    )

    // 特征向量化
    val to_vec = udf((arr:mutable.WrappedArray[Double])=>{
      Vectors.dense(arr.toArray)
    })

    val testVec = testDatumn.select('name,to_vec(array('job,'income,'age,'sex)) as "vec")

    // 加载训练好的模型
    val model = NaiveBayesModel.load("user_portrait/data/chugui/model")

    // 用模型来对待预测数据进行预测
    val res = model.transform(testVec)

    res.printSchema()
    res.show(100,false)

    val getProbFromVec = udf((vec:Vector)=>{
      vec(1)
    })

    res.select('name,getProbFromVec('probability) as "prob",'prediction).show(100,false)


    spark.close()


  }

}
